package org.honeynet.hbbackend.virustotal;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStreamWriter;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLConnection;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.Iterator;

import javax.annotation.PreDestroy;
import javax.annotation.Resource;
import javax.ejb.EJBException;
import javax.ejb.Schedule;
import javax.ejb.Singleton;
import javax.ejb.Startup;
import javax.ejb.TransactionAttribute;
import javax.ejb.TransactionAttributeType;
import javax.jms.ConnectionFactory;
import javax.jms.JMSException;
import javax.jms.Message;
import javax.jms.Session;
import javax.jms.Topic;
import javax.jms.TopicSubscriber;
import javax.sql.DataSource;

import org.codehaus.jackson.JsonNode;
import org.codehaus.jackson.map.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Singleton
@Startup
public class VirusTotalBinding {
	
	private final static String API_URL = "http://www.virustotal.com/api/get_file_report.json";
	private final static int RECV_TIMEOUT = 3000;
	
	private static Logger log = LoggerFactory.getLogger(VirusTotalBinding.class);
	
	@Resource(mappedName="jdbc/hbbackend")
	DataSource ds;
	
	@Resource(mappedName="jms/DurableConsumer/virustotal")
	private ConnectionFactory connectionFactory;
	
	@Resource(mappedName="jms/new_binary")
	private Topic dest;

	@Resource(name="apiKey")
	private String apiKey;
	
	private boolean stop = false;
	
	
	@Schedule(second="*/15", minute="*", hour="*", persistent=false)
	@TransactionAttribute(TransactionAttributeType.REQUIRES_NEW)
	public void work() {
		log.trace("work() called");
		
		javax.jms.Connection mq = null;
		Session sess = null;
		TopicSubscriber sub = null;
		
		try {
			mq = connectionFactory.createConnection();
			sess = mq.createSession(true, 0);
			sub = sess.createDurableSubscriber(dest, "virustotal");
			log.trace("created durable subscription");
			log.trace("starting receive loop");
			mq.start();
			
			Message jmsMsg;
			do {
				if (stop) {
					log.trace("encountered stop request: returning");
					return;
				}
				log.trace("receive");
				jmsMsg = sub.receive(RECV_TIMEOUT);
			} while (jmsMsg == null);
			
			//TODO message validation, check if type is unique?

			long id = jmsMsg.getLongProperty("id");
			String md5 = jmsMsg.getStringProperty("md5");
			log.debug("received new message id = {} md5 = {}", id, md5);

			store(id, retrieve(md5));
			return;
			
		} catch (JMSException e) {
			log.error("got JMSException", e);
			throw new EJBException(e);
		} finally {
			try {
//				if (sub != null) sub.close();
//				if (sess != null) sess.close();
				if (mq != null) mq.close();
			} catch (JMSException e) {
				log.error("got JMSException while trying to close", e);
				throw new EJBException(e);
			}
		}
		
	}
	
	
	private JsonNode retrieve(String md5) {
		log.trace("retrieve() called");

		String req =
			"resource" + "=" + md5 +
			"&" + "key" + "=" + apiKey;
		
		OutputStreamWriter out = null;
		InputStream in = null;
		JsonNode report;
		
		try {
			URL url = new URL(API_URL);
			URLConnection conn = url.openConnection();

			conn.setDoOutput(true);
			conn.setDoInput(true);
			
			out = new OutputStreamWriter(conn.getOutputStream());
			log.debug("sending request: " + req);
			out.write(req);
			out.flush();
			
			in = conn.getInputStream();
			ObjectMapper mapper = new ObjectMapper();
			report = mapper.readValue(in, JsonNode.class);
			
		} catch (MalformedURLException e) {
			log.error("got MalformedURLException", e);
			throw new EJBException(e);
		} catch (IOException e) {
			log.error("got IOException", e);
			throw new EJBException(e);
		} finally {
			try {
				if (out != null) out.close();
				if (in != null) in.close();
			} catch (IOException e) {
				log.error("got IOException", e);
				throw new EJBException(e);
			}
		}
		
		log.trace("json response: {}", report.toString());

		JsonNode resultField = report.findValue("result");
		
		if (resultField == null) {
			log.error("no result field in response");
			throw new EJBException("no result field in response");
		}
		
		int result = resultField.getIntValue();
		log.trace("result = {}", result);
		
		if (result == -1) {
			log.error("result = -1 (apiKey incorrect)");
			throw new EJBException("result = -1 (apiKey incorrect)");
		}
		
		if (result == -2) {
			log.error("result = -2 (exceeded request rate)");
			throw new EJBException("result = -2 (exceeded request rate)");
		}
		
		if (result != 1 && result != 0) {
			log.error("unkown result code = {}", result);
			throw new EJBException("unkown result code = " + result);
		}
		
		return report;
	}
	
	
	private void store(Long binaryId, JsonNode report) {
		log.trace("store() called");
		
		Connection db = null;
		PreparedStatement pStmt;
		ResultSet queryRes;

		int result;
		Timestamp timestamp;
		String permalink;
		Iterator<String> fieldNames;
		Iterator<JsonNode> fieldValues;
		int reportDataObjectSize;
		int zeroValueCount = 0;
		
		try {
			db = ds.getConnection();

			JsonNode resultField = report.findValue("result");
			result = resultField.getIntValue();

			if (result == 0) {
				pStmt = db.prepareStatement("insert into virustotal_reports (binary_id, resultcode) values (?,?)");
				pStmt.setLong(1, binaryId);
				pStmt.setInt(2, result);
				pStmt.executeUpdate();
				pStmt.close();
			} else {
				
				JsonNode reportArray = report.findValue("report");
				if (reportArray == null) {
					log.error("report array not found");
					throw new EJBException("report array not found");
				}
				log.trace("reportObject = {}", reportArray);
				
				JsonNode timeStampField = reportArray.get(0);
				if (timeStampField == null) {
					log.error("timestamp not found");
					throw new EJBException("timestamp not found");
				}
				try {
					timestamp = Timestamp.valueOf(timeStampField.getTextValue());
					log.trace("timestamp = {}", timestamp);
				} catch (NumberFormatException e) {
					log.error("couldn't parse timestamp: {}", timeStampField.getTextValue());
					throw new EJBException("couldn't parse timestamp: " + timeStampField.getTextValue());
				}
				
				JsonNode reportDataObject = reportArray.get(1);
				if (reportDataObject == null) {
					log.error("report data object not found");
					throw new EJBException("report data object not found");
				}
				log.trace("reportDataObject = {}", reportDataObject);
				
				// first, count zero values
				reportDataObjectSize = reportDataObject.size();
				fieldNames = reportDataObject.getFieldNames();
				fieldValues = reportDataObject.iterator();
				while (fieldValues.hasNext()) {
					String name = fieldNames.next();
					String value = fieldValues.next().getTextValue();
					
					if (name.equals("") && value.equals("")) {
						zeroValueCount++;
						continue;
					}
					if (value.equals("")) {
						zeroValueCount++;
					}
				}
				log.trace("zeroValueCount = {}, reportDataObjectSize = {}", zeroValueCount, reportDataObjectSize);
				
				JsonNode permalinkField = report.findValue("permalink");
				if (permalinkField == null) {
					log.error("permalink not found");
					throw new EJBException("permalink not found");
				}
				permalink = permalinkField.getTextValue();
				log.trace("permalink = {}", permalink);

				
				pStmt = db.prepareStatement("insert into virustotal_reports (binary_id, report_timestamp, resultcode, permalink) values (?,?,?,?) returning id");
				pStmt.setLong(1, binaryId);
				pStmt.setTimestamp(2, timestamp);
				pStmt.setInt(3, result);
				pStmt.setString(4, permalink);
				queryRes = pStmt.executeQuery();
				queryRes.next();
				long reportId = queryRes.getLong(1);
				log.debug("inserted into virustotal_reports (id = {} binary_id = {})", reportId, binaryId);
				queryRes.close();
				pStmt.close();
				
				if (zeroValueCount == reportDataObjectSize) {
					log.trace("empty report data, returning");
					return;
				}
				
				
				log.debug("inserting into virustotal_reports_data (report_id = {})", reportId);
				
				pStmt = db.prepareStatement("insert into virustotal_reports_data (report_id, engine, label) values (?,?,?)");
				
				fieldNames = reportDataObject.getFieldNames();
				fieldValues = reportDataObject.iterator();
				while (fieldValues.hasNext()) {
					String name = fieldNames.next();
					String value = fieldValues.next().getTextValue();
					
					if (name.equals("") && value.equals("")) {
						continue;
					}
					
					pStmt.setLong(1, reportId);
					pStmt.setString(2, name);
					pStmt.setString(3, value);
					pStmt.executeUpdate();
				}
				pStmt.close();
				
				return;
			}
		} catch (SQLException e) {
			log.error("got SQLException", e);
			throw new EJBException(e);
		} finally {
			try {
				if (db != null)	db.close();
			} catch (SQLException e) {
				log.error("got SQLException", e);
				throw new EJBException(e);
			}
		}
	}
	
	
	@PreDestroy
	public void stop() {
		stop = true;
	}
	
}
